{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import os\n",
    "import random\n",
    "import math\n",
    "import sys\n",
    "from PIL import Image\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">> Converting image 5798/5798\n",
      ">> Converting image 500/500\n",
      "生成tfcecord文件\n"
     ]
    }
   ],
   "source": [
    "#验证集数量\n",
    "_NUM_TEST = 500\n",
    "\n",
    "#随机种子\n",
    "_RANDOM_SEED = 0\n",
    "\n",
    "#数据集路径\n",
    "DATASET_DIR = \"D:/Tensorflow/captcha/images/\"\n",
    "\n",
    "#tfrecord文件存放路径\n",
    "TFRECORD_DIR = \"D:/Tensorflow/captcha/\"\n",
    "\n",
    "\n",
    "#判断tfrecord文件是否存在\n",
    "def _dataset_exists(dataset_dir):\n",
    "    for split_name in ['train', 'test']:\n",
    "        output_filename = os.path.join(dataset_dir,split_name + '.tfrecords')\n",
    "        if not tf.gfile.Exists(output_filename):\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "#获取所有验证码图片\n",
    "def _get_filenames_and_classes(dataset_dir):\n",
    "    photo_filenames = []\n",
    "    for filename in os.listdir(dataset_dir):\n",
    "        #获取文件路径\n",
    "        path = os.path.join(dataset_dir, filename)\n",
    "        photo_filenames.append(path)\n",
    "    return photo_filenames\n",
    "\n",
    "def int64_feature(values):\n",
    "    if not isinstance(values, (tuple, list)):\n",
    "        values = [values]\n",
    "    return tf.train.Feature(int64_list=tf.train.Int64List(value=values))\n",
    "\n",
    "def bytes_feature(values):\n",
    "    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))\n",
    "\n",
    "def image_to_tfexample(image_data, label0, label1, label2, label3):\n",
    "    #Abstract base class for protocol messages.\n",
    "    return tf.train.Example(features=tf.train.Features(feature={\n",
    "      'image': bytes_feature(image_data),\n",
    "      'label0': int64_feature(label0),\n",
    "      'label1': int64_feature(label1),\n",
    "      'label2': int64_feature(label2),\n",
    "      'label3': int64_feature(label3),\n",
    "    }))\n",
    "\n",
    "#把数据转为TFRecord格式\n",
    "def _convert_dataset(split_name, filenames, dataset_dir):\n",
    "    assert split_name in ['train', 'test']\n",
    "\n",
    "    with tf.Session() as sess:\n",
    "        #定义tfrecord文件的路径+名字\n",
    "        output_filename = os.path.join(TFRECORD_DIR,split_name + '.tfrecords')\n",
    "        with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:\n",
    "            for i,filename in enumerate(filenames):\n",
    "                try:\n",
    "                    sys.stdout.write('\\r>> Converting image %d/%d' % (i+1, len(filenames)))\n",
    "                    sys.stdout.flush()\n",
    "\n",
    "                    #读取图片\n",
    "                    image_data = Image.open(filename)  \n",
    "                    #根据模型的结构resize\n",
    "                    image_data = image_data.resize((224, 224))\n",
    "                    #灰度化\n",
    "                    image_data = np.array(image_data.convert('L'))\n",
    "                    #将图片转化为bytes\n",
    "                    image_data = image_data.tobytes()              \n",
    "\n",
    "                    #获取label\n",
    "                    labels = filename.split('/')[-1][0:4]\n",
    "                    num_labels = []\n",
    "                    for j in range(4):\n",
    "                        num_labels.append(int(labels[j]))\n",
    "                                            \n",
    "                    #生成protocol数据类型\n",
    "                    example = image_to_tfexample(image_data, num_labels[0], num_labels[1], num_labels[2], num_labels[3])\n",
    "                    tfrecord_writer.write(example.SerializeToString())\n",
    "                    \n",
    "                except IOError as e:\n",
    "                    print('Could not read:',filename)\n",
    "                    print('Error:',e)\n",
    "                    print('Skip it\\n')\n",
    "    sys.stdout.write('\\n')\n",
    "    sys.stdout.flush()\n",
    "\n",
    "#判断tfrecord文件是否存在\n",
    "if _dataset_exists(TFRECORD_DIR):\n",
    "    print('tfcecord文件已存在')\n",
    "else:\n",
    "    #获得所有图片\n",
    "    photo_filenames = _get_filenames_and_classes(DATASET_DIR)\n",
    "\n",
    "    #把数据切分为训练集和测试集,并打乱\n",
    "    random.seed(_RANDOM_SEED)\n",
    "    random.shuffle(photo_filenames)\n",
    "    training_filenames = photo_filenames[_NUM_TEST:]\n",
    "    testing_filenames = photo_filenames[:_NUM_TEST]\n",
    "\n",
    "    #数据转换\n",
    "    _convert_dataset('train', training_filenames, DATASET_DIR)\n",
    "    _convert_dataset('test', testing_filenames, DATASET_DIR)\n",
    "    print('生成tfcecord文件')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
